import torch
import torchvision
import torchvision.transforms as transforms

batch_size = 128
input_size = 32

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomCrop(32, 4),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

def getDataLoader(transforms=data_transforms):
    trainset = torchvision.datasets.CIFAR100(root='./data', train=True,
            download=False, transform=transforms['train'])
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR100(root='./data', train=False,
                download=False, transform=transforms['test'])
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                shuffle=False, num_workers=2)
    
    return (trainloader, testloader)

if __name__ == '__main__':
    trainloader, testloader = getDataLoader()
    print(len(trainloader), len(testloader))
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        print(inputs.size(), labels.size())
        break
    for i, data in enumerate(testloader, 0):
        inputs, labels = data
        print(inputs.size(), labels.size())
        break